import torch

# 简单的注意力机制如下：
# # ===========  定义Q,K,V 矩阵大小 ===============
# Q = torch.randn(3, 2)  # 3:批次  2：特征数
# K = torch.randn(8, 2)  # 8:实际Key的数量 2： 特征数
# V = torch.randn(8, 1)  # 8:实际value的数量  1: 输出特征值
#
# # ===========  点积注意力机制 ===============
# result = torch.softmax(Q @ K.T ,dim=-1)@ V
# print(result.shape)

# 复杂一点的注意力机制如下：
# Q,K,V 要求批次必须一致才能运算。
hidden_dims = 10
Q = torch.randn(8, 2, hidden_dims)  # 8:批次  2: 句子长度  10: 隐藏层的大小（w_q:(embedding_dims,hidden_dims)）
K = torch.randn(8, 4, hidden_dims)  # 8:实际Key的数量 4：句子长度
V = torch.randn(8, 4, hidden_dims)  # 8:实际value的数量  4：句子长度
# ===========  点积注意力机制 ===============
# 三维矩阵的点积运行，需要运算的维度是最后的两个位置
result = torch.softmax(Q @ K.transpose(1, 2), dim=-1) @ V
vocab_size = 5  # 词库大小为5
O = torch.randn(hidden_dims, vocab_size)
Z = result.reshape(-1, hidden_dims) @ O
print(Z.shape)
